from .Encode import Encode
from .Res_MCNN import Res_MCNN
from .GCN import GCN
import torch
import torch.nn as nn
class RF_STED_No_T(nn.Module):
    def __init__(self):
        super(RF_STED_No_T,self).__init__()
        self.gcn1 = GCN(in_features = 6,out_features = 1)
        self.gcn2 = GCN(in_features = 6,out_features = 1)
        self.gcn3 = GCN(in_features = 6,out_features = 1)
        self.gcn4 = GCN(in_features = 6,out_features = 1)
        self.encode = Encode()
        self.Res_CNN1 = Res_MCNN(num_convolutions=3, kernel_size=31, scale=2)
        self.Res_CNN2 = Res_MCNN(num_convolutions=3, kernel_size=33, scale=2)
        self.Res_CNN3 = Res_MCNN(num_convolutions=3, kernel_size=35, scale=2)   #num_convolutions越小越好  kernel_size居中  scale 2 差不多
        self.relu = nn.ReLU(inplace=False) #relu激活
    def forward(self,x,adj1,adj2,adj3,adj4):  #前向传播   x:[6,1,6,69,69]  adj:[4761,4761]
        #空间
        sp = x.squeeze(1)#(6,6,69,69)
        sp = sp.view((sp.shape[0],sp.shape[1],sp.shape[2]*sp.shape[3]))#(6,6,4761)
        sp = sp.permute(0,2,1)                           #torch.Size([6, 4761, 6])
        sp1 = self.gcn1(x = sp,adj = adj1).view(sp.shape[0],69,69,-1)     
        sp2 = self.gcn2(x = sp,adj = adj2).view(sp.shape[0],69,69,-1) 
        sp3 = self.gcn3(x = sp,adj = adj3).view(sp.shape[0],69,69,-1) 
        sp4 = self.gcn4(x = sp,adj = adj4).view(sp.shape[0],69,69,-1)    #[6，69,69，1]
        # 编码
        space = torch.cat((sp1.unsqueeze(-1), sp2.unsqueeze(-1), sp3.unsqueeze(-1), sp4.unsqueeze(-1)), dim=-1) #[6,69,69,1,4]
        space = space.view(space.shape[0], space.shape[1], space.shape[2] , space.shape[3] * space.shape[4])  #[6,69,69,4]
        encoder = space.sum(dim=3, keepdim=True)  #[6,69,69,1]
        encoder = sp2
        # 解码
        encoder = encoder.permute(0,3,1,2)   #[6,1,69,69]
        encoder = self.Res_CNN1(encoder) #[6,1,69,69]
        encoder = self.Res_CNN2(encoder) #[6,1,69,69]
        encoder = self.Res_CNN3(encoder) #[6,1,69,69]
        out = self.relu(encoder)
        return out
